{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "header-cell",
   "metadata": {},
   "source": [
    "# Qwen-Image-Edit 图像编辑演示\n",
    "\n",
    "本笔记本演示如何使用 Qwen-Image-Edit 模型进行各种图像编辑任务，包括：\n",
    "- 语义编辑（风格转换、视角变换等）\n",
    "- 外观编辑（添加/删除对象、修改颜色等）\n",
    "- 文本编辑（修改图像中的文字）\n",
    "- 链式编辑（逐步修正错误）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "install-deps",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 安装必要的依赖\n",
    "!pip install modelscope ipywidgets tqdm accelerate diffusers transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "check-versions",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 检查版本信息\n",
    "from importlib.metadata import version\n",
    "print(\"diffusers version:\", version(\"diffusers\"))\n",
    "print(\"torch version:\", version(\"torch\"))\n",
    "print(\"transformers version:\", version(\"transformers\"))\n",
    "print(\"PIL version:\", version(\"Pillow\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "download-model",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 下载 Qwen-Image-Edit 模型\n",
    "from modelscope import snapshot_download\n",
    "import os\n",
    "\n",
    "def download_edit_model(model_id=\"Qwen/Qwen-Image-Edit\", local_dir='./models/Qwen-Image-Edit'):\n",
    "    \"\"\"\n",
    "    下载图像编辑模型\n",
    "    \"\"\"\n",
    "    print(f\"开始下载 {model_id} 模型到: {local_dir}\")\n",
    "    \n",
    "    # 确保目录存在\n",
    "    os.makedirs(os.path.dirname(local_dir), exist_ok=True)\n",
    "    \n",
    "    # 下载模型\n",
    "    snapshot_download(\n",
    "        model_id, \n",
    "        local_dir=local_dir\n",
    "    )\n",
    "    \n",
    "    print(f\"模型下载完成，保存在: {local_dir}\")\n",
    "\n",
    "# 下载模型\n",
    "download_edit_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "setup-device",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置设备和数据类型\n",
    "import torch\n",
    "from PIL import Image\n",
    "import os\n",
    "\n",
    "# 检测设备\n",
    "if torch.cuda.is_available():\n",
    "    device = \"cuda\"\n",
    "    torch_dtype = torch.bfloat16\n",
    "    print(f\"使用 GPU: {torch.cuda.get_device_name()}\")\n",
    "    print(f\"GPU 内存: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB\")\n",
    "else:\n",
    "    device = \"cpu\"\n",
    "    torch_dtype = torch.float32\n",
    "    print(\"使用 CPU\")\n",
    "\n",
    "print(f\"设备: {device}, 数据类型: {torch_dtype}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "load-pipeline",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 加载图像编辑管道\n",
    "from diffusers import QwenImageEditPipeline\n",
    "\n",
    "model_path = \"./models/Qwen-Image-Edit\"\n",
    "\n",
    "print(\"正在加载 Qwen-Image-Edit 管道...\")\n",
    "pipeline = QwenImageEditPipeline.from_pretrained(model_path, torch_dtype=torch_dtype,use_safetensors=True,device_map=\"balanced\")\n",
    "print(\"管道加载完成\")\n",
    "\n",
    "# 设置数据类型和设备\n",
    "# pipeline.to(torch_dtype)\n",
    "# pipeline.to(device)\n",
    "# pipeline.set_progress_bar_config(disable=None)\n",
    "\n",
    "print(f\"管道已配置到 {device} 设备，数据类型: {torch_dtype}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "helper-functions",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 辅助函数\n",
    "def edit_image(image_path, prompt, output_path=None, seed=42, steps=50, cfg_scale=4.0):\n",
    "    \"\"\"\n",
    "    编辑图像的通用函数\n",
    "    \n",
    "    Args:\n",
    "        image_path: 输入图像路径\n",
    "        prompt: 编辑提示词\n",
    "        output_path: 输出图像路径（可选）\n",
    "        seed: 随机种子\n",
    "        steps: 推理步数\n",
    "        cfg_scale: CFG 缩放因子\n",
    "    \n",
    "    Returns:\n",
    "        PIL Image: 编辑后的图像\n",
    "    \"\"\"\n",
    "    # 加载输入图像\n",
    "    if isinstance(image_path, str):\n",
    "        image = Image.open(image_path).convert(\"RGB\")\n",
    "    else:\n",
    "        image = image_path\n",
    "    \n",
    "    print(f\"正在编辑图像...\")\n",
    "    print(f\"提示词: {prompt}\")\n",
    "    \n",
    "    # 编辑参数\n",
    "    inputs = {\n",
    "        \"image\": image,\n",
    "        \"prompt\": prompt,\n",
    "        \"generator\": torch.manual_seed(seed),\n",
    "        \"true_cfg_scale\": cfg_scale,\n",
    "        \"negative_prompt\": \" \",\n",
    "        \"num_inference_steps\": steps,\n",
    "    }\n",
    "    \n",
    "    # 执行编辑\n",
    "    with torch.inference_mode():\n",
    "        output = pipeline(**inputs)\n",
    "        edited_image = output.images[0]\n",
    "    \n",
    "    # 保存图像\n",
    "    if output_path:\n",
    "        edited_image.save(output_path)\n",
    "        print(f\"编辑后的图像已保存到: {os.path.abspath(output_path)}\")\n",
    "    \n",
    "    return edited_image\n",
    "\n",
    "def display_before_after(original_path, edited_image, title=\"图像编辑对比\"):\n",
    "    \"\"\"\n",
    "    显示编辑前后的对比图\n",
    "    \"\"\"\n",
    "    import matplotlib.pyplot as plt\n",
    "    \n",
    "    fig, axes = plt.subplots(1, 2, figsize=(12, 6))\n",
    "    \n",
    "    # 原图\n",
    "    if isinstance(original_path, str):\n",
    "        original = Image.open(original_path)\n",
    "    else:\n",
    "        original = original_path\n",
    "    axes[0].imshow(original)\n",
    "    axes[0].set_title(\"原图\")\n",
    "    axes[0].axis('off')\n",
    "    \n",
    "    # 编辑后\n",
    "    axes[1].imshow(edited_image)\n",
    "    axes[1].set_title(\"编辑后\")\n",
    "    axes[1].axis('off')\n",
    "    \n",
    "    plt.suptitle(title, fontsize=16)\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "print(\"辅助函数已定义完成\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "demo-section",
   "metadata": {},
   "source": [
    "## 演示案例\n",
    "\n",
    "以下是各种图像编辑任务的演示。请确保在 `./demo_images/` 目录下准备好相应的测试图像。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "create-demo-dir",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建演示目录\n",
    "demo_dir = \"./demo_images\"\n",
    "output_dir = \"./edited_images\"\n",
    "\n",
    "os.makedirs(demo_dir, exist_ok=True)\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "\n",
    "print(f\"演示目录已创建: {demo_dir}\")\n",
    "print(f\"输出目录已创建: {output_dir}\")\n",
    "print(\"\\n请将测试图像放入 demo_images 目录中\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "semantic-editing",
   "metadata": {},
   "source": [
    "### 1. 语义编辑演示"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "style-transfer",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 风格转换示例\n",
    "input_image = \"./demo_images/portrait.jpg\"  # 请准备一张人像图片\n",
    "\n",
    "if os.path.exists(input_image):\n",
    "    # 转换为宫崎骏风格\n",
    "    style_prompts = [\n",
    "        \"Transform this portrait into Studio Ghibli animation style, with soft colors and dreamy atmosphere\",\n",
    "        \"Convert to oil painting style with impressionist brushstrokes\",\n",
    "        \"Transform into cyberpunk style with neon lighting and futuristic elements\"\n",
    "    ]\n",
    "    \n",
    "    for i, prompt in enumerate(style_prompts):\n",
    "        output_path = f\"{output_dir}/style_transfer_{i+1}.png\"\n",
    "        edited = edit_image(input_image, prompt, output_path, seed=42+i)\n",
    "        display_before_after(input_image, edited, f\"风格转换 {i+1}: {prompt[:30]}...\")\n",
    "else:\n",
    "    print(f\"请将人像图片保存为: {input_image}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "viewpoint-change",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 视角变换示例\n",
    "input_image = \"./demo_images/object.jpg\"  # 请准备一张物体图片\n",
    "\n",
    "if os.path.exists(input_image):\n",
    "    viewpoint_prompts = [\n",
    "        \"Rotate this object 90 degrees to show the side view\",\n",
    "        \"Show the back view of this object by rotating it 180 degrees\",\n",
    "        \"Change the viewing angle to show this object from above\"\n",
    "    ]\n",
    "    \n",
    "    for i, prompt in enumerate(viewpoint_prompts):\n",
    "        output_path = f\"{output_dir}/viewpoint_{i+1}.png\"\n",
    "        edited = edit_image(input_image, prompt, output_path, seed=100+i)\n",
    "        display_before_after(input_image, edited, f\"视角变换 {i+1}\")\n",
    "else:\n",
    "    print(f\"请将物体图片保存为: {input_image}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "appearance-editing",
   "metadata": {},
   "source": [
    "### 2. 外观编辑演示"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "object-addition",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 添加对象示例\n",
    "input_image = \"./demo_images/scene.jpg\"  # 请准备一张场景图片\n",
    "\n",
    "if os.path.exists(input_image):\n",
    "    addition_prompts = [\n",
    "        \"Add a red signboard with text 'OPEN' to this scene, include realistic reflection\",\n",
    "        \"Add a cute cat sitting in the foreground\",\n",
    "        \"Add some colorful balloons floating in the sky\"\n",
    "    ]\n",
    "    \n",
    "    for i, prompt in enumerate(addition_prompts):\n",
    "        output_path = f\"{output_dir}/addition_{i+1}.png\"\n",
    "        edited = edit_image(input_image, prompt, output_path, seed=200+i)\n",
    "        display_before_after(input_image, edited, f\"添加对象 {i+1}\")\n",
    "else:\n",
    "    print(f\"请将场景图片保存为: {input_image}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "color-change",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 颜色修改示例\n",
    "input_image = \"./demo_images/colorful_object.jpg\"  # 请准备一张彩色物体图片\n",
    "\n",
    "if os.path.exists(input_image):\n",
    "    color_prompts = [\n",
    "        \"Change the main object's color to bright purple\",\n",
    "        \"Make this image black and white except for one red element\",\n",
    "        \"Change the background to a warm sunset color palette\"\n",
    "    ]\n",
    "    \n",
    "    for i, prompt in enumerate(color_prompts):\n",
    "        output_path = f\"{output_dir}/color_change_{i+1}.png\"\n",
    "        edited = edit_image(input_image, prompt, output_path, seed=400+i)\n",
    "        display_before_after(input_image, edited, f\"颜色修改 {i+1}\")\n",
    "else:\n",
    "    print(f\"请将彩色物体图片保存为: {input_image}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "text-editing",
   "metadata": {},
   "source": [
    "### 3. 文本编辑演示"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "text-modification",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 文本修改示例\n",
    "input_image = \"./demo_images/text_image.jpg\"  # 请准备一张包含文字的图片\n",
    "\n",
    "if os.path.exists(input_image):\n",
    "    text_prompts = [\n",
    "        \"Change the text in this image to 'HELLO WORLD' in the same style\",\n",
    "        \"Replace any English text with Chinese characters '欢迎光临'\",\n",
    "        \"Change the color of the text to bright blue while keeping everything else the same\"\n",
    "    ]\n",
    "    \n",
    "    for i, prompt in enumerate(text_prompts):\n",
    "        output_path = f\"{output_dir}/text_edit_{i+1}.png\"\n",
    "        edited = edit_image(input_image, prompt, output_path, seed=500+i)\n",
    "        display_before_after(input_image, edited, f\"文本编辑 {i+1}\")\n",
    "else:\n",
    "    print(f\"请将包含文字的图片保存为: {input_image}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "chained-editing",
   "metadata": {},
   "source": [
    "### 4. 链式编辑演示"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "progressive-editing",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 链式编辑示例 - 逐步完善图像\n",
    "input_image = \"./demo_images/base_image.jpg\"  # 请准备一张基础图片\n",
    "\n",
    "if os.path.exists(input_image):\n",
    "    # 第一步编辑\n",
    "    step1_prompt = \"Add a beautiful garden with flowers in the background\"\n",
    "    step1_result = edit_image(input_image, step1_prompt, f\"{output_dir}/chain_step1.png\", seed=600)\n",
    "    display_before_after(input_image, step1_result, \"链式编辑 - 步骤1: 添加花园\")\n",
    "    \n",
    "    # 第二步编辑（基于第一步结果）\n",
    "    step2_prompt = \"Add a wooden bench in the garden for people to sit\"\n",
    "    step2_result = edit_image(step1_result, step2_prompt, f\"{output_dir}/chain_step2.png\", seed=601)\n",
    "    display_before_after(step1_result, step2_result, \"链式编辑 - 步骤2: 添加长椅\")\n",
    "    \n",
    "    # 第三步编辑（基于第二步结果）\n",
    "    step3_prompt = \"Add warm golden hour lighting to create a romantic atmosphere\"\n",
    "    step3_result = edit_image(step2_result, step3_prompt, f\"{output_dir}/chain_step3.png\", seed=602)\n",
    "    display_before_after(step2_result, step3_result, \"链式编辑 - 步骤3: 添加黄金时光照明\")\n",
    "    \n",
    "    # 显示最终对比\n",
    "    display_before_after(input_image, step3_result, \"链式编辑 - 最终对比\")\n",
    "    \n",
    "else:\n",
    "    print(f\"请将基础图片保存为: {input_image}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "batch-processing",
   "metadata": {},
   "source": [
    "### 5. 批量处理演示"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "batch-edit",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 批量处理示例\n",
    "import glob\n",
    "\n",
    "# 获取所有演示图片\n",
    "demo_images = glob.glob(f\"{demo_dir}/*.jpg\") + glob.glob(f\"{demo_dir}/*.png\")\n",
    "\n",
    "if demo_images:\n",
    "    batch_prompt = \"Enhance this image with better lighting and more vibrant colors\"\n",
    "    \n",
    "    print(f\"找到 {len(demo_images)} 张图片进行批量处理\")\n",
    "    \n",
    "    for i, img_path in enumerate(demo_images[:3]):  # 限制处理前3张图片\n",
    "        print(f\"\\n处理第 {i+1} 张图片: {os.path.basename(img_path)}\")\n",
    "        \n",
    "        output_path = f\"{output_dir}/batch_enhanced_{i+1}.png\"\n",
    "        edited = edit_image(img_path, batch_prompt, output_path, seed=700+i)\n",
    "        \n",
    "        # 显示对比\n",
    "        display_before_after(img_path, edited, f\"批量增强 {i+1}\")\n",
    "        \n",
    "    print(\"\\n批量处理完成！\")\n",
    "else:\n",
    "    print(\"未找到演示图片，请在 demo_images 目录中添加图片\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "summary",
   "metadata": {},
   "source": [
    "## 总结\n",
    "\n",
    "本笔记本演示了 Qwen-Image-Edit 模型的强大图像编辑能力：\n",
    "\n",
    "1. **语义编辑**: 风格转换、视角变换等高级编辑\n",
    "2. **外观编辑**: 添加/删除对象、修改颜色等精确编辑\n",
    "3. **文本编辑**: 修改图像中的文字内容和样式\n",
    "4. **链式编辑**: 通过多步编辑实现复杂的图像变换\n",
    "5. **批量处理**: 高效处理多张图像\n",
    "\n",
    "### 使用建议：\n",
    "- 使用清晰、具体的提示词获得更好的编辑效果\n",
    "- 对于复杂编辑，可以使用链式编辑方法逐步完善\n",
    "- 调整 `cfg_scale` 和 `num_inference_steps` 参数来平衡质量和速度\n",
    "- 使用不同的随机种子获得多样化的编辑结果"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
